package com.study.flink.sink

import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.{RichSinkFunction, SinkFunction}
import org.apache.flink.streaming.api.scala.StreamExecutionEnvironment
import org.apache.hadoop.hbase.client.{Connection, ConnectionFactory, Put, Table}
import org.apache.hadoop.hbase.util.Bytes
import org.apache.hadoop.hbase.{HBaseConfiguration, HConstants, TableName}

/**
  * HBase Sink Demo
  *
  * @author stephen
  * @date 2019-07-23 18:02
  */
object ScalaHBaseSinkDemo {

  def main(args: Array[String]): Unit = {

    // 1 获取执行环境
    val env = StreamExecutionEnvironment.getExecutionEnvironment

    // 2 获取输入数据
    import org.apache.flink.api.scala._
    // john,20,male,1001
    // lucy,22,female,1002
    val dataStream = env.socketTextStream("localhost", 9999)

    // 3 Transformation
    val resultStream = dataStream.flatMap(_.split(" "))
      .map(line => {
        val arr = line.split(",")
        Student(arr(0), arr(1).toInt, arr(2), arr(3))
      })

    // 4 输出
    val hbaseSink = new HBaseSink("student", "info")
    resultStream.addSink(hbaseSink)

    // 5 启动任务
    env.execute("HBase Sink Demo")
  }
}

case class Student(name: String, age: Int, sex: String, sid: String)

class HBaseSink(tableName: String, family: String) extends RichSinkFunction[Student] {

  var conn: Connection = _
  val ZK_HOST = "localhost:2181"

  override def open(parameters: Configuration): Unit = {
    super.open(parameters)
    val conf = HBaseConfiguration.create()
    conf.set(HConstants.ZOOKEEPER_QUORUM, ZK_HOST)
    conn = ConnectionFactory.createConnection(conf)
  }

  override def invoke(student: Student, context: SinkFunction.Context[_]): Unit = {
    super.invoke(student, context)

    val t: Table = conn.getTable(TableName.valueOf(tableName))
    val put: Put = new Put(Bytes.toBytes(student.sid))
    put.addColumn(Bytes.toBytes(family), Bytes.toBytes("name"), Bytes.toBytes(student.name))
    put.addColumn(Bytes.toBytes(family), Bytes.toBytes("age"), Bytes.toBytes(student.age))
    put.addColumn(Bytes.toBytes(family), Bytes.toBytes("sex"), Bytes.toBytes(student.sex))
    t.put(put)
    t.close()
  }

  override def close(): Unit = {
    super.close()
    conn.close()
  }
}
